{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.agents import AgentExecutor, create_tool_calling_agent\n",
    "from langchain_core.prompts import ChatPromptTemplate\n",
    "from ChatGLM_new import tongyi_llm,zhipu_llm\n",
    "from langchain_core.tools import tool\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "def magic_function(input: int) -> int:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    return input + 2\n",
    "\n",
    "\n",
    "tools = [magic_function]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{'input': 'what is the value of magic_function(3)?',\n",
       " 'output': 'The value of magic_function(3) is 5.'}"
      ]
     },
     "execution_count": 6,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "prompt = ChatPromptTemplate.from_messages(\n",
    "    [\n",
    "        (\"system\", \"You are a helpful assistant\"),\n",
    "        (\"human\", \"{input}\"),\n",
    "        # Placeholders fill up a **list** of messages\n",
    "        (\"placeholder\", \"{agent_scratchpad}\"),\n",
    "    ]\n",
    ")\n",
    "\n",
    "query = \"what is the value of magic_function(3)?\"\n",
    "\n",
    "agent = create_tool_calling_agent(zhipu_llm, tools, prompt)\n",
    "\n",
    "\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools)\n",
    "\n",
    "agent_executor.invoke({\"input\": query})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "AgentExecutor(agent=RunnableMultiActionAgent(runnable=RunnableAssign(mapper={\n",
       "  agent_scratchpad: RunnableLambda(lambda x: format_to_tool_messages(x['intermediate_steps']))\n",
       "})\n",
       "| ChatPromptTemplate(input_variables=['input'], input_types={'agent_scratchpad': typing.List[typing.Union[langchain_core.messages.ai.AIMessage, langchain_core.messages.human.HumanMessage, langchain_core.messages.chat.ChatMessage, langchain_core.messages.system.SystemMessage, langchain_core.messages.function.FunctionMessage, langchain_core.messages.tool.ToolMessage]]}, partial_variables={'agent_scratchpad': []}, messages=[SystemMessagePromptTemplate(prompt=PromptTemplate(input_variables=[], template='You are a helpful assistant')), HumanMessagePromptTemplate(prompt=PromptTemplate(input_variables=['input'], template='{input}')), MessagesPlaceholder(variable_name='agent_scratchpad', optional=True)])\n",
       "| RunnableBinding(bound=ChatOpenAI(verbose=True, client=<openai.resources.chat.completions.Completions object at 0x000001EAC4521C50>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x000001EAD2D0A690>, model_name='glm-4', openai_api_key=SecretStr('**********'), openai_api_base='https://open.bigmodel.cn/api/paas/v4', openai_proxy=''), kwargs={'tools': [{'type': 'function', 'function': {'name': 'magic_function', 'description': 'Applies a magic function to an input.', 'parameters': {'type': 'object', 'properties': {'input': {'type': 'integer'}}, 'required': ['input']}}}]})\n",
       "| ToolsAgentOutputParser(), input_keys_arg=[], return_keys_arg=[], stream_runnable=True), tools=[StructuredTool(name='magic_function', description='Applies a magic function to an input.', args_schema=<class 'pydantic.v1.main.magic_functionSchema'>, func=<function magic_function at 0x000001EAC1BB05E0>)])"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "agent_executor"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#! pip install langgraph"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "@tool\n",
    "def magic_function(input) -> int:\n",
    "    \"\"\"Applies a magic function to an input.\"\"\"\n",
    "    print(\"Applying magic function to input:\", input)\n",
    "    return input * 2\n",
    "\n",
    "\n",
    "tools = [magic_function]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "tool_names = [i.name for i in tools]\n",
    "\n",
    "\n",
    "template = '''Answer the following questions as best you can. You have access to the following tools:\n",
    "\n",
    "{tools}\n",
    "\n",
    "Use the following format:\n",
    "\n",
    "Question: the input question you must answer\n",
    "Thought: you should always think about what to do\n",
    "Action: the action to take, should be one of [{tool_names}]\n",
    "Action Input: the input to the action\n",
    "Observation: the result of the action\n",
    "... (this Thought/Action/Action Input/Observation can repeat N times)\n",
    "Thought: I now know the final answer\n",
    "Final Answer: the final answer to the original input question\n",
    "\n",
    "Begin!\n",
    "\n",
    "Question: {input}\n",
    "Thought:{agent_scratchpad}'''\n",
    "\n",
    "prompt = PromptTemplate.from_template(template)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Applying magic function to input: 3\n",
      "Observation\n",
      "Applying magic function to input: 3\n",
      "\n",
      "Observation\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'input': 'what is the value of magic_function(3)?',\n",
       " 'output': 'The value of magic_function(3) is 3.'}"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from langchain.agents import AgentExecutor,create_react_agent\n",
    "\n",
    "agent  = create_react_agent(zhipu_llm, tools, prompt)\n",
    "\n",
    "agent_executor = AgentExecutor(agent=agent, tools=tools)\n",
    "\n",
    "\n",
    "agent_executor.invoke({\"input\": \"what is the value of magic_function(3)?\"})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 31,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "RunnableAssign(mapper={\n",
       "  agent_scratchpad: RunnableLambda(lambda x: format_log_to_str(x['intermediate_steps']))\n",
       "})\n",
       "| PromptTemplate(input_variables=['agent_scratchpad', 'input'], partial_variables={'tools': 'magic_function(input: int) -> int - Applies a magic function to an input.', 'tool_names': 'magic_function'}, template='Answer the following questions as best you can. You have access to the following tools:\\n\\n{tools}\\n\\nUse the following format:\\n\\nQuestion: the input question you must answer\\nThought: you should always think about what to do\\nAction: the action to take, should be one of [{tool_names}]\\nAction Input: the input to the action\\nObservation: the result of the action\\n... (this Thought/Action/Action Input/Observation can repeat N times)\\nThought: I now know the final answer\\nFinal Answer: the final answer to the original input question\\n\\nBegin!\\n\\nQuestion: {input}\\nThought:{agent_scratchpad}')\n",
       "| RunnableBinding(bound=ChatOpenAI(verbose=True, client=<openai.resources.chat.completions.Completions object at 0x000001EAC4521C50>, async_client=<openai.resources.chat.completions.AsyncCompletions object at 0x000001EAD2D0A690>, model_name='glm-4', openai_api_key=SecretStr('**********'), openai_api_base='https://open.bigmodel.cn/api/paas/v4', openai_proxy=''), kwargs={'stop': ['\\nObservation']})\n",
       "| ReActSingleInputOutputParser()"
      ]
     },
     "execution_count": 31,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "app"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
